#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Piper机械臂Joycon手柄控制系统
作者: Assistant
日期: 2024-12-19
功能: 通过Joycon手柄实时控制Piper机械臂的6自由度移动和旋转
版本: 1.0

主要功能:
1. 加速度计控制X-Y-Z轴移动（需要按住ZL/ZR按键）
2. 陀螺仪控制RX-RY-RZ轴旋转（需要按住ZL/ZR按键）
3. ZL/ZR按键作为控制开关
4. 实时状态显示和传感器数据监控
5. 安全初始位姿和重置功能
"""

import math
import os
import sys
import time
import threading
import select
import termios
import tty
from typing import Dict, List, Tuple, Optional
from dataclasses import dataclass

# 添加项目路径
sys.path.append(os.path.join(os.path.dirname(__file__), ".."))

import mujoco
from mujoco import minimize
import numpy as np
from piper_sdk import *
from scipy.spatial.transform import Rotation

# 尝试导入 pyjoycon 库
try:
    from pyjoycon import JoyCon, get_R_id, get_L_id
    PYOYCON_AVAILABLE = True
except ImportError:
    print("警告: pyjoycon 库未安装，将使用 pygame 作为备选方案")
    print("安装命令: pip install joycon-python")
    PYOYCON_AVAILABLE = False

# 尝试导入 pygame 库
try:
    import pygame
    PYGAME_AVAILABLE = True
except ImportError:
    print("警告: pygame 库未安装")
    print("安装命令: pip install pygame")
    PYGAME_AVAILABLE = False

@dataclass
class GyroscopeData:
    """陀螺仪数据结构"""
    x: float  # 绕X轴角速度 (度/秒)
    y: float  # 绕Y轴角速度 (度/秒)
    z: float  # 绕Z轴角速度 (度/秒)
    magnitude: float  # 角速度大小

@dataclass
class AccelerometerData:
    """加速度计数据结构"""
    x: float  # X轴加速度 (m/s²)
    y: float  # Y轴加速度 (m/s²)
    z: float  # Z轴加速度 (m/s²)
    magnitude: float  # 加速度大小
    pitch: float  # 俯仰角 (度)
    roll: float  # 横滚角 (度)

class JoyconPiperController:
    """
    Joycon手柄控制Piper机械臂类
    支持6自由度实时控制（位置+旋转）和夹爪操作
    """

    def __init__(self, piper):
        """
        初始化Joycon控制器

        Args:
            piper: Piper机械臂接口对象
        """
        self.piper = piper

        # 初始化Joycon手柄
        self._init_joycon_controller()

        # 机械臂控制参数
        self.factor = 1000  # 单位转换因子（毫米转微米，度转微度）
        self.arm_speed = 50  # 机械臂速度（cm/s）
        self.move_speed = 20  # 移动速度（厘米/秒）
        self.rotation_speed = 30  # 旋转速度（度/秒）
        self.control_frequency = 60.0  # 控制频率（Hz）
        self.move_distance_per_frame = self.move_speed / self.control_frequency  # 每帧移动距离（厘米）
        self.rotation_angle_per_frame = self.rotation_speed / self.control_frequency  # 每帧旋转角度（度）

        # 当前位置和初始位置
        self.current_position = None  # 当前机械臂位置 [X, Y, Z, RX, RY, RZ]
        self.initial_position = None  # 初始位置
        self.safe_initial_pose = [57.0, 0.0, 215.0, 0.0, 85.0, 0.0]  # 安全初始位姿

        # 传感器控制参数
        self.accel_deadzone = 0.5  # 加速度计死区 (m/s²)
        self.gyro_deadzone = 50.0  # 陀螺仪死区 (度/秒)
        self.accel_sensitivity = 0.1  # 加速度计灵敏度
        self.gyro_sensitivity = 0.01  # 陀螺仪灵敏度

        # 控制开关状态
        self.zl_pressed = False  # ZL按键状态
        self.zr_pressed = False  # ZR按键状态
        self.control_enabled = False  # 控制是否启用

        # 传感器数据
        self.gyro_data_L = GyroscopeData(0.0, 0.0, 0.0, 0.0)
        self.gyro_data_R = GyroscopeData(0.0, 0.0, 0.0, 0.0)
        self.accel_data_L = AccelerometerData(0.0, 0.0, 0.0, 0.0, 0.0, 0.0)
        self.accel_data_R = AccelerometerData(0.0, 0.0, 0.0, 0.0, 0.0, 0.0)

        # 运动平滑参数
        self.smoothness_factor = 0.8  # 运动平滑因子
        self.last_target_position = None  # 上一次目标位置
        self.motion_smoothing_enabled = True  # 运动平滑开关

        # 性能监控参数
        self.frame_count = 0
        self.last_fps_check = time.time()
        self.actual_fps = 0

        # 键盘输入处理
        self.keyboard_thread = None
        self.running = True

        # 初始化机械臂
        self._initialize_robot()

        # 添加IK求解器
        self.ik = MujocoIKSolver()

    def _init_joycon_controller(self):
        """初始化Joycon手柄"""
        self.joycon_L = None
        self.joycon_R = None
        self.pygame_joysticks = []
        
        # 初始化 pyjoycon 手柄
        if PYOYCON_AVAILABLE:
            self._init_pyjoycon()
        
        # 初始化 pygame 作为备选方案
        if PYGAME_AVAILABLE and (not PYOYCON_AVAILABLE or (self.joycon_L is None and self.joycon_R is None)):
            self._init_pygame()
        
        if self.joycon_L is None and self.joycon_R is None and not self.pygame_joysticks:
            raise RuntimeError("未检测到Joycon手柄，请确保手柄已连接")

    def _init_pyjoycon(self):
        """初始化 pyjoycon 手柄"""
        print("正在初始化 pyjoycon 手柄...")
        
        try:
            # 获取 Joycon 设备 ID
            joycon_id_L = get_L_id()
            joycon_id_R = get_R_id()
            
            # 连接左 Joycon
            if joycon_id_L and None not in joycon_id_L:
                try:
                    self.joycon_L = JoyCon(*joycon_id_L)
                    print(f"✅ 左 Joycon 连接成功")
                except Exception as e:
                    print(f"❌ 左 Joycon 连接失败: {e}")
                    self.joycon_L = None
            
            # 连接右 Joycon
            if joycon_id_R and None not in joycon_id_R:
                try:
                    self.joycon_R = JoyCon(*joycon_id_R)
                    print(f"✅ 右 Joycon 连接成功")
                except Exception as e:
                    print(f"❌ 右 Joycon 连接失败: {e}")
                    self.joycon_R = None
            
            if self.joycon_L is None and self.joycon_R is None:
                print("❌ 未检测到任何 Joycon 设备")
                return False
                
            return True
            
        except Exception as e:
            print(f"❌ pyjoycon 初始化失败: {e}")
            return False

    def _init_pygame(self):
        """初始化 pygame 作为备选方案"""
        print("正在初始化 pygame 手柄...")
        
        try:
            pygame.init()
            pygame.joystick.init()
            
            joystick_count = pygame.joystick.get_count()
            print(f"检测到 {joystick_count} 个手柄设备")
            
            if joystick_count == 0:
                print("❌ 未检测到任何手柄设备")
                return False
            
            # 初始化所有手柄
            for i in range(joystick_count):
                joystick = pygame.joystick.Joystick(i)
                joystick.init()
                
                name = joystick.get_name()
                print(f"手柄 {i}: {name}")
                print(f"  轴数量: {joystick.get_numaxes()}")
                print(f"  按钮数量: {joystick.get_numbuttons()}")
                print(f"  帽子数量: {joystick.get_numhats()}")
                
                self.pygame_joysticks.append(joystick)
            
            return True
            
        except Exception as e:
            print(f"❌ pygame 初始化失败: {e}")
            return False

    def _initialize_robot(self):
        """初始化机械臂"""
        print("正在初始化机械臂...")

        try:
            # 连接机械臂
            print("步骤1: 连接机械臂...")
            self.piper.ConnectPort()
            print("机械臂连接成功")

            # 使能机械臂
            print("步骤2: 使能机械臂...")
            self.piper.EnableArm(7)
            print("机械臂使能命令已发送")

            # 使能检测
            print("步骤3: 检测使能状态...")
            self._enable_robot()

            # 移动到安全初始位姿
            print("步骤4: 移动到安全初始位姿...")
            self._move_to_initial_pose()

            # 等待移动完成并获取实际位置
            print("步骤5: 获取当前位置...")
            time.sleep(2)
            self.current_position = self._get_current_position()

            # 如果获取位置失败，使用预定义位置
            if self.current_position is None:
                print("警告: 无法获取当前位置，使用预定义位姿")
                self.current_position = self.safe_initial_pose.copy()
            else:
                print("成功获取当前位置")

            # 保存初始位置用于重置功能
            self.initial_position = self.current_position.copy()
            self.last_target_position = self.current_position.copy()
            print("初始位置已保存")

            # 显示初始位置信息
            print(
                f"初始位置: X={self.current_position[0]:.1f}, Y={self.current_position[1]:.1f}, Z={self.current_position[2]:.1f}"
            )
            print(
                f"         RX={self.current_position[3]:.1f}, RY={self.current_position[4]:.1f}, RZ={self.current_position[5]:.1f}"
            )
            print(f"控制参数: 移动速度={self.move_speed}cm/s, 旋转速度={self.rotation_speed}°/s")
            print(f"          控制频率={self.control_frequency}Hz")
            print(f"          加速度计死区={self.accel_deadzone}m/s², 陀螺仪死区={self.gyro_deadzone}°/s")
            print(f"          加速度计灵敏度={self.accel_sensitivity}, 陀螺仪灵敏度={self.gyro_sensitivity}")
            print("机械臂初始化完成！")

        except Exception as e:
            print(f"机械臂初始化失败: {e}")
            import traceback
            traceback.print_exc()
            raise

    def _enable_robot(self):
        """使能机械臂并检测使能状态"""
        enable_flag = False
        timeout = 5  # 超时时间（秒）
        start_time = time.time()

        while not enable_flag:
            elapsed_time = time.time() - start_time
            print("检查使能状态...")

            try:
                # 获取所有电机的使能状态
                low_spd_info = self.piper.GetArmLowSpdInfoMsgs()
                enable_flag = (
                    low_spd_info.motor_1.foc_status.driver_enable_status
                    and low_spd_info.motor_2.foc_status.driver_enable_status
                    and low_spd_info.motor_3.foc_status.driver_enable_status
                    and low_spd_info.motor_4.foc_status.driver_enable_status
                    and low_spd_info.motor_5.foc_status.driver_enable_status
                    and low_spd_info.motor_6.foc_status.driver_enable_status
                )
            except Exception as e:
                print(f"获取使能状态失败: {e}")
                enable_flag = False

            # 如果未使能，尝试重新使能
            if not enable_flag:
                self.piper.EnableArm(7)
                self.piper.GripperCtrl(0, 1000, 0x01, 0)

            # 超时检查
            if elapsed_time > timeout:
                print("程序自动使能超时,退出程序")
                exit(0)

            time.sleep(1)
        print("机械臂使能成功")

    def _move_to_initial_pose(self):
        """移动到安全初始位姿"""
        try:
            # 将位置转换为微米单位
            X = round(self.safe_initial_pose[0] * self.factor)
            Y = round(self.safe_initial_pose[1] * self.factor)
            Z = round(self.safe_initial_pose[2] * self.factor)
            RX = round(self.safe_initial_pose[3] * self.factor)
            RY = round(self.safe_initial_pose[4] * self.factor)
            RZ = round(self.safe_initial_pose[5] * self.factor)

            print(f"移动到初始位姿: X={X}, Y={Y}, Z={Z}, RX={RX}, RY={RY}, RZ={RZ}")

            # 使用慢速移动到初始位置（速度15，确保安全）
            self.piper.MotionCtrl_2(0x01, 0x00, 15, 0x00)
            self.piper.EndPoseCtrl(X, Y, Z, RX, RY, RZ)
            time.sleep(3)  # 等待移动完成

        except Exception as e:
            print(f"移动到初始位姿失败: {e}")

    def _wait_motion_done(self):
        """等待机械臂运动完成"""
        start_t = time.time()
        timeout = 5  # 超时时间（秒）

        while self._is_in_motion():
            time.sleep(0.01)  # 10ms检查间隔
            if time.time() - start_t > timeout:
                print("等待运动完成超时")
                return False
        return True

    def _is_in_motion(self):
        """检查机械臂是否在运动"""
        try:
            return bool(self.piper.GetArmStatus().arm_status.motion_status)
        except:
            return False

    def _get_current_position(self):
        """获取当前机械臂末端位姿"""
        try:
            # 使用正确的API调用方法获取位姿信息
            endpose = self.piper.GetArmEndPoseMsgs().end_pose

            # 转换为毫米和度
            position = [
                endpose.X_axis / self.factor,  # 微米转毫米
                endpose.Y_axis / self.factor,
                endpose.Z_axis / self.factor,
                endpose.RX_axis / self.factor,  # 微度转度
                endpose.RY_axis / self.factor,
                endpose.RZ_axis / self.factor,
            ]
            return position

        except Exception as e:
            print(f"获取位姿失败: {e}")
            return None

    def IKmove(self, X, Y, Z, RX, RY, RZ, speed=40):
        """
        X, Y, Z, RX, RY, RZ, J0 are machine values (int) of End effector
        """
        print(f"IKmoveJoint: X={X}, Y={Y}, Z={Z}, RX={RX}, RY={RY}, RZ={RZ}, speed={speed}")
        pos = np.array([X, Y, Z]) / 1e6
        rpy = np.deg2rad(np.array([RX, RY, RZ]) / 1000)
        # get current joints:
        arm_status = self.piper.GetArmJointMsgs().joint_state
        jo = [
            arm_status.joint_1,
            arm_status.joint_2,
            arm_status.joint_3,
            arm_status.joint_4,
            arm_status.joint_5,
            arm_status.joint_6,
            0,  # gripper 1
            0,
        ]
        j0 = np.deg2rad(np.asarray(jo) / 1000)
        j = self.ik.solve_ik_xyz_rpy(pos=pos, rpy=rpy, j0=j0)
        machine_joints = (np.rad2deg(j[:6]) * 1000).astype(int)
        self.piper.ModeCtrl(0x01, 0x01, speed, 0x00)
        self.piper.JointCtrl(*machine_joints)

    def _apply_motion_smoothing(self, target_pos):
        """
        应用运动平滑算法

        Args:
            target_pos: 目标位置

        Returns:
            smoothed_pos: 平滑后的位置
        """
        if not self.motion_smoothing_enabled or self.last_target_position is None:
            return target_pos

        # 位置平滑
        smoothed_pos = []
        for i in range(6):
            current_val = self.last_target_position[i]
            target_val = target_pos[i]
            smoothed_val = current_val + (target_val - current_val) * (1 - self.smoothness_factor)
            smoothed_pos.append(smoothed_val)

        return smoothed_pos

    def _move_robot_to_position(self, new_pos, wait_completion=False):
        """移动机械臂到指定位置"""
        try:
            # 只有在需要等待完成时才等待（如重置位置时）
            if wait_completion:
                self._wait_motion_done()

            # 应用运动平滑
            if self.motion_smoothing_enabled:
                new_pos = self._apply_motion_smoothing(new_pos)

            # 转换为微米单位
            X = round(new_pos[0] * self.factor)
            Y = round(new_pos[1] * self.factor)
            Z = round(new_pos[2] * self.factor)
            RX = round(new_pos[3] * self.factor)
            RY = round(new_pos[4] * self.factor)
            RZ = round(new_pos[5] * self.factor)

            # 发送位置控制命令
            self.IKmove(X, Y, Z, RX, RY, RZ, speed=self.arm_speed)

            # 更新位置记录
            self.current_position = new_pos.copy()
            self.last_target_position = new_pos.copy()

        except Exception as e:
            print(f"移动失败: {e}")

    def _calculate_magnitude(self, x: float, y: float, z: float) -> float:
        """计算三维向量的模长"""
        return math.sqrt(x*x + y*y + z*z)

    def _calculate_angles(self, x: float, y: float, z: float) -> Tuple[float, float]:
        """计算俯仰角和横滚角"""
        # 俯仰角 (pitch) - 绕X轴旋转
        pitch = math.atan2(y, math.sqrt(x*x + z*z)) * 180 / math.pi
        
        # 横滚角 (roll) - 绕Y轴旋转
        roll = math.atan2(-x, math.sqrt(y*y + z*z)) * 180 / math.pi
        
        return pitch, roll

    def _update_sensor_data(self):
        """更新传感器数据"""
        # 更新左 Joycon 传感器数据
        if self.joycon_L:
            try:
                # 陀螺仪数据
                gyro_x = self.joycon_L.get_gyro_x()
                gyro_y = self.joycon_L.get_gyro_y()
                gyro_z = self.joycon_L.get_gyro_z()
                gyro_mag = self._calculate_magnitude(gyro_x, gyro_y, gyro_z)
                self.gyro_data_L = GyroscopeData(gyro_x, gyro_y, gyro_z, gyro_mag)
                
                # 加速度计数据
                accel_x = self.joycon_L.get_accel_x()
                accel_y = self.joycon_L.get_accel_y()
                accel_z = self.joycon_L.get_accel_z()
                accel_mag = self._calculate_magnitude(accel_x, accel_y, accel_z)
                pitch, roll = self._calculate_angles(accel_x, accel_y, accel_z)
                self.accel_data_L = AccelerometerData(accel_x, accel_y, accel_z, accel_mag, pitch, roll)
                
            except Exception as e:
                print(f"左 Joycon 传感器数据读取失败: {e}")
        
        # 更新右 Joycon 传感器数据
        if self.joycon_R:
            try:
                # 陀螺仪数据
                gyro_x = self.joycon_R.get_gyro_x()
                gyro_y = self.joycon_R.get_gyro_y()
                gyro_z = self.joycon_R.get_gyro_z()
                gyro_mag = self._calculate_magnitude(gyro_x, gyro_y, gyro_z)
                self.gyro_data_R = GyroscopeData(gyro_x, gyro_y, gyro_z, gyro_mag)
                
                # 加速度计数据
                accel_x = self.joycon_R.get_accel_x()
                accel_y = self.joycon_R.get_accel_y()
                accel_z = self.joycon_R.get_accel_z()
                accel_mag = self._calculate_magnitude(accel_x, accel_y, accel_z)
                pitch, roll = self._calculate_angles(accel_x, accel_y, accel_z)
                self.accel_data_R = AccelerometerData(accel_x, accel_y, accel_z, accel_mag, pitch, roll)
                
            except Exception as e:
                print(f"右 Joycon 传感器数据读取失败: {e}")

    def _read_joycon_inputs(self):
        """读取Joycon手柄输入"""
        # 处理 pygame 事件
        if PYGAME_AVAILABLE and self.pygame_joysticks:
            for event in pygame.event.get():
                if event.type == pygame.QUIT:
                    return False

        # 更新传感器数据
        self._update_sensor_data()
        
        # 读取 pyjoycon 手柄数据
        if self.joycon_L:
            try:
                status_L = self.joycon_L.get_status()
                self._process_pyjoycon_status(status_L, "L")
            except Exception as e:
                print(f"左 Joycon 状态读取失败: {e}")
        
        if self.joycon_R:
            try:
                status_R = self.joycon_R.get_status()
                self._process_pyjoycon_status(status_R, "R")
            except Exception as e:
                print(f"右 Joycon 状态读取失败: {e}")

        return True

    def _process_pyjoycon_status(self, status: Dict, side: str):
        """处理 pyjoycon 状态数据"""
        # 处理按钮数据
        buttons = status.get('buttons', {})
        
        # 检查ZL/ZR按键状态
        if side == "L":
            left_buttons = buttons.get('left', {})
            self.zl_pressed = left_buttons.get('zl', False)
        elif side == "R":
            right_buttons = buttons.get('right', {})
            self.zr_pressed = right_buttons.get('zr', False)
        
        # 检查Home按键状态（用于重置功能）
        shared_buttons = buttons.get('shared', {})
        home_pressed = shared_buttons.get('home', False)
        
        # 如果Home按键被按下，重置到初始位置
        if home_pressed:
            self.reset_to_initial_position()

    def _handle_sensor_control(self):
        """处理传感器控制"""
        # 检查控制开关状态
        self.control_enabled = self.zl_pressed or self.zr_pressed
        
        if not self.control_enabled:
            return

        # 使用Joycon的传感器数据进行控制
        # 优先使用左Joycon，如果没有则使用右Joycon
        accel_x, accel_y, accel_z = 0, 0, 0
        gyro_x, gyro_y, gyro_z = 0, 0, 0
        
        if self.joycon_L:
            # 使用左Joycon的传感器数据
            accel_x = self.accel_data_L.x
            accel_y = self.accel_data_L.y
            accel_z = self.accel_data_L.z
            
            gyro_x = self.gyro_data_L.x
            gyro_y = self.gyro_data_L.y
            gyro_z = self.gyro_data_L.z
        elif self.joycon_R:
            # 如果没有左Joycon，使用右Joycon的传感器数据
            accel_x = self.accel_data_R.x
            accel_y = self.accel_data_R.y
            accel_z = self.accel_data_R.z
            
            gyro_x = self.gyro_data_R.x
            gyro_y = self.gyro_data_R.y
            gyro_z = self.gyro_data_R.z
            
            # 应用死区
            if abs(accel_x) < self.accel_deadzone:
                accel_x = 0
            if abs(accel_y) < self.accel_deadzone:
                accel_y = 0
            if abs(accel_z - 9.8) < self.accel_deadzone:  # 重力补偿
                accel_z = 0
            else:
                accel_z = accel_z - 9.8  # 减去重力加速度
                
            if abs(gyro_x) < self.gyro_deadzone:
                gyro_x = 0
            if abs(gyro_y) < self.gyro_deadzone:
                gyro_y = 0
            if abs(gyro_z) < self.gyro_deadzone:
                gyro_z = 0
            
            # 计算控制量
            if self.current_position is not None:
                new_pos = self.current_position.copy()
                
                # 位置控制（加速度计）
                if abs(accel_x) > 0:
                    new_pos[0] += accel_x * self.accel_sensitivity * self.move_distance_per_frame * 10
                if abs(accel_y) > 0:
                    new_pos[1] += accel_y * self.accel_sensitivity * self.move_distance_per_frame * 10
                if abs(accel_z) > 0:
                    new_pos[2] += accel_z * self.accel_sensitivity * self.move_distance_per_frame * 10
                
                # 旋转控制（陀螺仪）
                if abs(gyro_x) > 0:
                    new_pos[3] += gyro_x * self.gyro_sensitivity * self.rotation_angle_per_frame
                if abs(gyro_y) > 0:
                    new_pos[4] += gyro_y * self.gyro_sensitivity * self.rotation_angle_per_frame
                if abs(gyro_z) > 0:
                    new_pos[5] += gyro_z * self.gyro_sensitivity * self.rotation_angle_per_frame
                
                # 移动机械臂
                self._move_robot_to_position(new_pos, wait_completion=False)

    def reset_to_initial_position(self):
        """重置到初始位置"""
        print("正在重置到初始位置...")
        if self.initial_position is None:
            print("错误: 初始位置未设置，无法重置")
            return

        print(
            f"目标位置: X={self.initial_position[0]:.1f}, Y={self.initial_position[1]:.1f}, Z={self.initial_position[2]:.1f}"
        )
        print(
            f"         RX={self.initial_position[3]:.1f}, RY={self.initial_position[4]:.1f}, RZ={self.initial_position[5]:.1f}"
        )

        try:
            self._move_robot_to_position(self.initial_position, wait_completion=True)
            print("已重置到初始位置")
        except Exception as e:
            print(f"重置位置失败: {e}")
            import traceback
            traceback.print_exc()

    def open_gripper(self):
        """打开夹爪（100mm宽度）"""
        print("正在打开夹爪...")
        try:
            # 等待运动完成
            self._wait_motion_done()
            # 发送夹爪控制命令 - 100mm宽度
            self.piper.GripperCtrl(100 * 1000, 1000, 0x01, 0)
            # 等待夹爪操作完成
            time.sleep(0.3)
            print("夹爪已打开")
        except Exception as e:
            print(f"打开夹爪失败: {e}")

    def close_gripper(self):
        """关闭夹爪（完全闭合）"""
        print("正在关闭夹爪...")
        try:
            # 等待运动完成
            self._wait_motion_done()
            # 发送夹爪控制命令 - 0mm宽度（完全闭合）
            self.piper.GripperCtrl(0, 1000, 0x01, 0)
            # 等待夹爪操作完成
            time.sleep(0.3)
            print("夹爪已关闭")
        except Exception as e:
            print(f"关闭夹爪失败: {e}")

    def _keyboard_input_handler(self):
        """键盘输入处理线程"""
        # 保存终端设置
        old_settings = termios.tcgetattr(sys.stdin)
        try:
            # 设置终端为非阻塞模式
            tty.setraw(sys.stdin.fileno())
            
            while self.running:
                # 检查是否有键盘输入
                if select.select([sys.stdin], [], [], 0.1)[0]:
                    key = sys.stdin.read(1)
                    if key == '\x1b':  # ESC键
                        print("\n检测到ESC键，正在退出...")
                        self.running = False
                        break
                    elif key == 'h' or key == 'H':  # Home键
                        print("\n检测到H键，重置到初始位置...")
                        self.reset_to_initial_position()
                    elif key == 'g' or key == 'G':  # 夹爪控制
                        print("\n检测到G键，切换夹爪状态...")
                        # 这里可以添加夹爪切换逻辑
                        
        except Exception as e:
            print(f"键盘输入处理错误: {e}")
        finally:
            # 恢复终端设置
            termios.tcsetattr(sys.stdin, termios.TCSADRAIN, old_settings)

    def update_performance_monitor(self):
        """性能监控"""
        self.frame_count += 1
        current_time = time.time()

        if current_time - self.last_fps_check >= 1.0:  # 每秒更新一次
            self.actual_fps = self.frame_count / (current_time - self.last_fps_check)
            self.frame_count = 0
            self.last_fps_check = current_time

    def run(self):
        """运行主控制循环"""
        print("\n=============== Piper机械臂Joycon控制系统 ===============")
        print("作者: Assistant")
        print("日期: 2024-12-19")
        print("\n控制说明:")
        print("  [传感器控制 - 需要按住ZL或ZR按键]")
        print("  - 加速度计X轴: 控制机械臂X轴前后移动")
        print("  - 加速度计Y轴: 控制机械臂Y轴左右移动")
        print("  - 加速度计Z轴: 控制机械臂Z轴上下移动")
        print("  - 陀螺仪X轴: 控制机械臂RX轴旋转")
        print("  - 陀螺仪Y轴: 控制机械臂RY轴旋转")
        print("  - 陀螺仪Z轴: 控制机械臂RZ轴旋转")
        print("  [控制开关]")
        print("  - ZL按键: 启用左Joycon控制")
        print("  - ZR按键: 启用右Joycon控制")
        print("  - 松开按键: 停止控制")
        print("  [其他功能]")
        print("  - Home按键: 重置到初始位置")
        print("  - ESC键: 退出程序")
        print("  - H键: 重置到初始位置")
        print("  - G键: 切换夹爪状态（待实现）")
        print("\n控制参数:")
        print(f"  - 移动速度: {self.move_speed}cm/s")
        print(f"  - 旋转速度: {self.rotation_speed}°/s")
        print(f"  - 控制频率: {self.control_frequency}Hz")
        print(f"  - 加速度计死区: {self.accel_deadzone}m/s²")
        print(f"  - 陀螺仪死区: {self.gyro_deadzone}°/s")
        print(f"  - 加速度计灵敏度: {self.accel_sensitivity}")
        print(f"  - 陀螺仪灵敏度: {self.gyro_sensitivity}")
        print("=" * 65)

        try:
            # 启动键盘输入处理线程
            self.keyboard_thread = threading.Thread(target=self._keyboard_input_handler, daemon=True)
            self.keyboard_thread.start()
            
            # 主控制循环 - 60Hz频率
            control_interval = 1.0 / self.control_frequency  # 约16.67ms
            last_control_time = time.time()

            while self.running:
                current_time = time.time()

                # 确保60Hz控制频率
                if current_time - last_control_time >= control_interval:
                    # 读取手柄输入
                    if not self._read_joycon_inputs():
                        break

                    # 处理传感器控制
                    self._handle_sensor_control()

                    # 性能监控
                    self.update_performance_monitor()

                    # 显示实时状态信息
                    if self.current_position is not None:
                        control_status = "🔴 控制中" if self.control_enabled else "⚪ 待机"
                        zl_status = "🔴 ZL" if self.zl_pressed else "⚪ ZL"
                        zr_status = "🔴 ZR" if self.zr_pressed else "⚪ ZR"
                        
                        print(
                            f"\r当前位置: X={self.current_position[0]:6.1f}, Y={self.current_position[1]:6.1f}, Z={self.current_position[2]:6.1f} | "
                            f"RX={self.current_position[3]:6.1f}, RY={self.current_position[4]:6.1f}, RZ={self.current_position[5]:6.1f} | "
                            f"{control_status} | {zl_status} {zr_status} | FPS: {self.actual_fps:4.1f}Hz",
                            end="",
                            flush=True,
                        )

                    last_control_time = current_time
                else:
                    # 短暂休眠以节省CPU
                    time.sleep(0.001)  # 1ms

        except KeyboardInterrupt:
            print("\n\n程序被用户中断")
        except Exception as e:
            print(f"\n\n程序运行出错: {e}")
            import traceback
            traceback.print_exc()
        finally:
            # 清理资源
            print("\n正在清理资源...")
            try:
                # 停止键盘输入线程
                self.running = False
                if self.keyboard_thread and self.keyboard_thread.is_alive():
                    self.keyboard_thread.join(timeout=1.0)
                
                # 断开Joycon手柄
                if self.joycon_L:
                    self.joycon_L.disconnect_device()
                if self.joycon_R:
                    self.joycon_R.disconnect_device()
                if PYGAME_AVAILABLE:
                    pygame.quit()
                print("Joycon手柄已断开连接")
            except:
                pass
            print("程序已退出")


class MujocoIKSolver:
    """
    The radius is the weight of the orientation residual (in Lie algebra)
    The smaller the radius, the more the orientation residual is ignored.

    Forward functions in mujoco:

    mujoco.mj_kinematics(self.model, self.data):
        compute the kinematics of the model and update the data inplace

    mujoco.mj_forward(self.model, self.data)
        = {
        1. mj_kinematics(model, data),  # update the position and the rotations
        2. compute the constraint forces and the accelerations
    }

    mujoco.mj_step(self.model, self.data):
        = {
        1. mj_forward(model, data),  # update the position and the rotations, forces and accelerations
        2. Do the integration of the dynamics
    }

    In this case, only position and rotation are updated, so mj_kinematics is enough.

    Mujoco use quaternion as [w, x, y, z]
    pinocchio use quaternion as [x, y, z, w]
    """

    def __init__(self, radius=0.04, reg=0.01):
        self.model = mujoco.MjModel.from_xml_path("piper_description/mujoco_model/piper_description.xml")
        self.data = mujoco.MjData(self.model)

        # Reset the state to the "zero_position" keyframe.
        self.go_zero()

        # Bounds at the joint limits.
        self.bounds = [self.model.jnt_range[:, 0], self.model.jnt_range[:, 1]]

        # Inital guess is the 'home' keyframe.
        self.j0 = self.model.key("home").qpos
        self.radius = radius
        self.reg = reg

    def go_zero(self):
        key = mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_KEY, "zero_position")
        mujoco.mj_resetDataKeyframe(self.model, self.data, key)
        mujoco.mj_forward(self.model, self.data)

    def move_target_to(self, pos, quat):
        self.data.mocap_pos[self.model.body("target").mocapid] = pos
        self.data.mocap_quat[self.model.body("target").mocapid] = quat
        mujoco.mj_kinematics(self.model, self.data)

    def ik_residual(self, joint, pos=None, quat=None, reg_target=None):
        """Residual for inverse kinematics.

        Args:
        joint: joint angles.
        pos: target position for the end effector.
        quat: target orientation for the end effector.
        radius: scaling of the 3D cross.

        Returns:
        The residual of the Inverse Kinematics task.
        """
        # Move the mocap body to the target
        tar_id = self.model.body("target").mocapid
        self.data.mocap_pos[tar_id] = self.model.body("target").pos if pos is None else pos
        self.data.mocap_quat[tar_id] = self.model.body("target").quat if quat is None else quat

        # Set qpos, compute forward kinematics.
        res = []
        for i in range(joint.shape[1]):
            self.data.qpos = joint[:, i]
            mujoco.mj_kinematics(self.model, self.data)

            # Position residual: can also use Lie algebra to compute the residual
            res_pos = self.data.site("effector").xpos - self.data.site("target").xpos

            # Effector quat, use mju_mat2quat.
            effector_quat = np.empty(4)
            mujoco.mju_mat2Quat(effector_quat, self.data.site("effector").xmat)

            # Target quat, exploit the fact that the site is aligned with the body.
            target_quat = self.data.body("target").xquat

            # Orientation residual: quaternion difference.
            res_quat = np.empty(3)
            mujoco.mju_subQuat(res_quat, target_quat, effector_quat)
            res_quat *= self.radius

            # Regularization residual. To avoid the joint angles from being too far from the home position.
            # reg_target=x_prev is to regularize the joint angles to be close to the previous joint angles
            reg_target = self.model.key("home").qpos if reg_target is None else reg_target
            res_reg = self.reg * (joint[:, i] - reg_target)

            res_i = np.hstack((res_pos, res_quat, res_reg))
            res.append(np.atleast_2d(res_i).T)

        return np.hstack(res)

    def ik_jac(self, joint, res, pos=None, quat=None):
        """Analytic Jacobian of inverse kinematics residual

        Args:
        joint: joint angles.
        pos: target position for the end effector.
        quat: target orientation for the end effector.
        radius: scaling of the 3D cross.

        Returns:
        The Jacobian of the Inverse Kinematics task.
        """
        # least_squares() passes the value of the residual at x which is sometimes
        # useful, but we don't need it here.
        del res

        # Call mj_kinematics and mj_comPos (required for Jacobians).
        mujoco.mj_kinematics(self.model, self.data)
        mujoco.mj_comPos(self.model, self.data)

        # Get end-effector site Jacobian.
        jac_pos = np.empty((3, self.model.nv))
        jac_quat = np.empty((3, self.model.nv))
        mujoco.mj_jacSite(self.model, self.data, jac_pos, jac_quat, self.data.site("effector").id)

        # Get Deffector, the 3x3 mju_subquat Jacobian
        effector_quat = np.empty(4)
        mujoco.mju_mat2Quat(effector_quat, self.data.site("effector").xmat)
        target_quat = self.data.body("target").xquat
        Deffector = np.empty((3, 3))
        # compute the jacobian: partial(target_quat - effector_quat) / partial(Deffector_quat)
        mujoco.mjd_subQuat(target_quat, effector_quat, None, Deffector)

        # Rotate into target frame, multiply by subQuat Jacobian, scale by radius.
        target_mat = self.data.site("target").xmat.reshape(3, 3)
        mat = self.radius * Deffector.T @ target_mat.T
        jac_quat = mat @ jac_quat

        # Regularization Jacobian.
        jac_reg = self.reg * np.eye(self.model.nv)

        return np.vstack((jac_pos, jac_quat, jac_reg))

    def solve_ik(self, pos=None, quat=None, j0=None, x_prev=None):
        """
        Solve the IK problem.
        pos: target position for the end effector.
        quat: target orientation for the end effector.
        j0: initial guess of the joint angles. If None, use the home position.
        """
        if j0 is None:
            # initial guess is the home position
            j0 = self.j0
        ik_target = lambda joint: self.ik_residual(joint, pos=pos, quat=quat, reg_target=x_prev)
        jac_target = lambda joint, r: self.ik_jac(joint, r, pos=pos, quat=quat)

        joint, _ = minimize.least_squares(
            j0, ik_target, self.bounds, jacobian=jac_target, verbose=0, check_derivatives=False, max_iter=1000
        )
        return joint

    def solve_ik_xyz_rpy(self, pos=None, rpy=None, j0=None, x_prev=None):
        """
        Solve the IK problem.
        pos: target position for the end effector.
        rpy: row, pitch, yaw of the end effector (in radians)
        j0: initial guess of the joint angles. If None, use the home position.
        """
        xyz_quat = Rotation.from_euler("xyz", rpy).as_quat()  # (x, y, z, w)
        quat = np.array([xyz_quat[3], xyz_quat[0], xyz_quat[1], xyz_quat[2]])  # (w, x, y, z)
        return self.solve_ik(pos=pos, quat=quat, j0=j0, x_prev=x_prev)


def main():
    """主函数"""
    try:
        # 创建Piper机械臂对象
        piper = C_PiperInterface("can0")

        # 创建Joycon控制器
        controller = JoyconPiperController(piper)

        # 运行控制循环
        controller.run()

    except Exception as e:
        print(f"程序启动失败: {e}")
        import traceback
        traceback.print_exc()


if __name__ == "__main__":
    main()
